"""
将每一条微调数据尽可能拆分成较少的数据量
schema 中不需要指令
跟据标注数据准备 event 数据集（数据中包括 ent 和 event，但不包括 relation）
注意事项：
    # fix1: 是否需要去掉文本首尾的空格？
    # fix2(DONE): 将 ent 的 type 属性的连字符从 '_' 改为 '-'
    # fix3: 是否需要过滤跨界线的关系？
    # fix4(DONE): 将 event 的所有 args 的 role 转为小写
    for arg in item['args']:
        arg['role'] = arg['role'].lower()
    # fix5(DONE): 为每个 event 的 event-type 加上 "action-" 前缀
    item["event_type"] = "action-" + item["event_type"]
    # fix6(DONE): 在划定一条数据所涉及的句子的从开始到结束的范围后，相应地要调整区间坐标（减去 offset），使左端点为 0，注意使用深拷贝！！！
    # fix7(DONE): 验证数据集拆分的正确性
    # fix8: 增加负向数据，即没有 event 的数据
"""

import os
import json
import random
import copy
from nltk.tokenize import sent_tokenize

cur_dir = os.path.dirname(os.path.abspath(__file__))

### 断句
# 每一句话拆分一下，作为一条微调数据

### 数据格式
class DataTemplate:
    def __init__(
            self, 
            id, 
            ent_list, 
            rel_list, 
            event_dict, 
            ans_ent_list, 
            ans_rel_list, 
            ans_event_list, 
            text, 
            bg
        ):
        self.id = id
        self.ent_list = ent_list
        self.rel_list = rel_list
        self.event_dict = event_dict
        self.ans_ent_list = ans_ent_list
        self.ans_rel_list = ans_rel_list
        self.ans_event_list = ans_event_list
        self.text = text
        self.bg = bg

    def assemble_data(self):
        return {
            "id": self.id, 
            "schema": {
                "ent": self.ent_list, 
                "rel": self.rel_list, 
                "event": self.event_dict
            }, 
            "ans": {
                "ent": self.ans_ent_list, 
                "rel": self.ans_rel_list, 
                "event": self.ans_event_list
            }, 
            "text": self.text, 
            "bg": self.bg
        }

# region 预处理数据集的一系列方法
# 从 .txt 文件中读取故事文本
def get_story_from_txt(cur_dir: str, filename: str):
    try:
        with open(os.path.join(cur_dir, filename), 'r', encoding='utf-8') as f:
            story_en = f.read()
        return story_en
    except:
        print(f"Error: {filename} for story not found in {cur_dir}")
        return None

# generate story sentence list of positions
def get_story_span_list(story_en):
    sentences = sent_tokenize(story_en)
    start_pos = 0
    end_pos = 0
    story_span_list = []
    for sentence in sentences:
        end_pos = start_pos + len(sentence)
        template_dict = { "span": [start_pos, end_pos], "sentence": sentence }
        story_span_list.append(template_dict)
        start_pos = end_pos + 1 # because of a space exists between sentences
    return story_span_list
# 生成一条数据的 id
def gen_id(data_type: str, id_num: int):
    return f"story.event.{data_type}.{id_num}"

# 从文件中读取 schema 的 ent_list
def get_ent_list_from_file(dir_path: str, filename: str):
    try:
        with open(os.path.join(dir_path, filename), "r", encoding="utf-8") as f:
            ent_list = json.load(f)
            return ent_list
    except:
        print(f"Error: {filename} for ent_list not found in {dir_path}")
        return None

# 从文件中读取 schema 的 rel_list
def get_rel_list_from_file():
    pass

# 从文件中读取 schema 的 event_dict
def get_event_dict_from_file(dir_path: str, filename: str):
    try:
        with open(os.path.join(dir_path, filename), "r", encoding="utf-8") as f:
            event_dict = json.load(f)
            return event_dict
    except:
        print(f"Error: {filename} for event_dict not found in {dir_path}")
        return None
    
# 计算 event 的 trigger 和所有 args 列表元素能够一起覆盖到的最大的 span 范围
def get_cover_span(temp_event_list: list, story_span_list: list):
    start_index = 9999999
    end_index = 0
    for event in temp_event_list:
        if event["trigger"]["span"][0] < start_index:
            start_index = event["trigger"]["span"][0]
        if event["trigger"]["span"][1] > end_index:
            end_index = event["trigger"]["span"][1]
        for arg in event["args"]:
            if arg["span"][0] < start_index:
                start_index = arg["span"][0]
            if arg["span"][1] > end_index:
                end_index = arg["span"][1]
    # 从 start_index 和 end_index 的字符开始，分别找到 temp_event_list 的第一句话的开始下标和最后一句话的结束下标
    # there must has a space between sentences for processed data, so it can preprocess the scale of each sentence after split sentences with nltk
    for story_span in story_span_list: # update start and end index
        if story_span['span'][0] <= start_index and start_index <= story_span['span'][1]:
            start_index = story_span['span'][0]
        if story_span['span'][0] <= end_index and end_index <= story_span['span'][1]:
            end_index = story_span['span'][1]
    return start_index, end_index

# 从标注数据 json 文件中提取 train、test、dev 数据的 ans_ent_list、ans_rel_list、ans_event_list、text
def gen_ans_ent_rel_event_text_from_json_file(cur_dir: str, filename: str, story_en: str, story_span_list: list):
    # 从文件中读取所有的 ent、event
    with open(os.path.join(cur_dir, filename), "r", encoding="utf-8") as f:
        data = json.load(f)
        ent_list = data["ans"]["ent"]
        event_list = data["ans"]["action"]
    # 跟据每一个 event 找到对应的故事文本 text 和包含的实体 ent，然后添加到 total_data 列表中
    total_data = []
    temp_event_list = []
    event_num = random.randint(1, 3) # 使用随机数 1/2/3 确定一条数据中包含的 event 数量
    cnt = 0 # 计数器，用于记录当前已经积累的 event 数量
    for event in event_list:
        if cnt < event_num:
            temp_event_list.append(event)
            cnt += 1
        else:
            # 跟据 temp_event_list 得到对应的实体列表和事件列表，并添加到 ent_lists 和 event_lists 中
            pinch_ent_list = []
            pinch_event_list = []
            for a_event in temp_event_list:
                # 跟据 a_event 中 args 列表中每个元素的 span 找到每个 ent，并添加到 pinch_ent_list 中
                for a_ent in a_event["args"]:
                    for ent in ent_list:
                        if a_ent["span"][0] == ent["span"][0] and a_ent["span"][1] == ent["span"][1]:
                            a_ent['role'] = a_ent['role'].lower() # fix4: 将 event 的所有 args 的 role 转为小写
                            ent['type'] = ent['type'].replace('_', '-') # fix2: 将 ent 的 type 属性的连字符从 '_' 改为 '-'
                            pinch_ent_list.append(copy.deepcopy(ent)) # 深拷贝，防止相同元素被一起修改
                            break
                # 向 pinch_event_list 中添加 a_event
                a_event["event_type"] = "action-" + a_event["event_type"] # fix5: 为每个 event 的 event-type 加上 "action-" 前缀
                pinch_event_list.append(copy.deepcopy(a_event)) # 深拷贝，防止相同元素被一起修改
            # 计算 trigger 和所有 args 列表元素能够一起覆盖到的最大的 span 范围
            text_start, text_end = get_cover_span(temp_event_list, story_span_list)
            # 跟据 span 范围提取所有涉及到的句子，作为 text，添加到 text_list 中
            ans_text = story_en[text_start:text_end]
            # 用 DataTemplate 组装一条数据
            a_piece_of_data = DataTemplate(
                id="", # 在 shuffle 的时候赋值
                ent_list=get_ent_list_from_file(cur_dir, "ent_list.json"),
                rel_list=[],
                event_dict=get_event_dict_from_file(cur_dir, "event_dict.json"),
                ans_ent_list=pinch_ent_list,
                ans_rel_list=[],
                ans_event_list=pinch_event_list,
                text=ans_text,
                bg=""
            )
            # 处理 span 的偏移量
            # 处理 ans_ent_list
            # print("haha")
            for a_ent in a_piece_of_data.ans_ent_list:
                a_ent["span"][0] -= text_start
                a_ent["span"][1] -= text_start
            # 处理 ans_event_list
            for a_event in a_piece_of_data.ans_event_list:
                a_event["trigger"]["span"][0] -= text_start
                a_event["trigger"]["span"][1] -= text_start
                for a_arg in a_event["args"]:
                    a_arg["span"][0] -= text_start
                    a_arg["span"][1] -= text_start
            # 将 a_piece_of_data 添加到 total_data 中
            total_data.append(a_piece_of_data)
            # 更新变量
            temp_event_list = []
            cnt = 0
            event_num = random.randint(1, 3)

    # 按比例划分数据集
    total_data_len = len(total_data)
    split_ratios = {
        'train': 0.8,
        'test': 0.1,
        'dev': 0.1
    }
    train_dataset = []
    test_dataset = []
    dev_dataset = []
    for i in range(total_data_len):
        if i < total_data_len * split_ratios['train']:
            train_dataset.append(total_data[i])
        elif i < total_data_len * (split_ratios['train'] + split_ratios['test']):
            test_dataset.append(total_data[i])
        else:
            dev_dataset.append(total_data[i])

    return total_data, train_dataset, test_dataset, dev_dataset

# 生成一条数据的 bg
def gen_bg():
    pass

# 验证数据集的正确性
def verify_dataset(dataset_list: list):
    for dataset in dataset_list:
        for data in dataset:
            # 验证 ans_ent_list
            for data_ent in data.ans_ent_list:
                original_text = data.text[data_ent["span"][0]:data_ent["span"][1]]
                extracted_text = data_ent["text"]
                if original_text != extracted_text:
                    print(f"id: {data.id} 的 ans_ent_list 验证失败，{original_text} != {extracted_text}")
                    return False
            # 验证 ans_event_list
            for data_event in data.ans_event_list:
                # 验证 trigger
                original_text = data.text[data_event["trigger"]["span"][0]:data_event["trigger"]["span"][1]]
                extracted_text = data_event["trigger"]["text"]
                if original_text != extracted_text:
                    print(f"id: {data.id} 的 ans_event_list 中 trigger 验证失败，{original_text} != {extracted_text}")
                    return False
                # 验证 args
                for data_arg in data_event["args"]:
                    original_text = data.text[data_arg["span"][0]:data_arg["span"][1]]
                    extracted_text = data_arg["text"]
                    if original_text != extracted_text:
                        print(f"id: {data.id} 的 ans_event_list 中 args 验证失败，{original_text} != {extracted_text}")
                        return False
    return True

# 生成负向数据，方法是添加删去数据特征的数据，当然重复使用会造成数据冗余，需要避免
def gen_neg_data(total_data: list, train_dataset: list, test_dataset: list, dev_dataset: list):
    neg_total_data = []
    # 删除每条数据的 text 中的实体和事件，并清空 ans_ent_list 和 ans_event_list
    for data in total_data:
        # 标记 text 的哪些位置被删除，删除的位置标记为 1，初始化为长度为 data.text 的全为 0 的列表
        remove_flag = [0] * len(data.text)
        if data.ans_ent_list or data.ans_event_list:
            for data_ent in data.ans_ent_list:
                remove_flag[data_ent["span"][0]:data_ent["span"][1]] = [1] * (data_ent["span"][1] - data_ent["span"][0])
            for data_event in data.ans_event_list:
                remove_flag[data_event["trigger"]["span"][0]:data_event["trigger"]["span"][1]] = [1] * (data_event["trigger"]["span"][1] - data_event["trigger"]["span"][0])
                for data_arg in data_event["args"]:
                    remove_flag[data_arg["span"][0]:data_arg["span"][1]] = [1] * (data_arg["span"][1] - data_arg["span"][0])
            # 删除 text 中的实体和事件
            neg_text = "".join([data.text[i] for i in range(len(data.text)) if remove_flag[i] == 0]).replace("  ", " ")
            # 构造 a_piece_of_neg_data
            a_piece_of_neg_data = DataTemplate(
                id="", # 在 shuffle 的时候赋值
                ent_list=get_ent_list_from_file(cur_dir, "ent_list.json"),
                rel_list=[],
                event_dict=get_event_dict_from_file(cur_dir, "event_dict.json"),
                ans_ent_list=[],
                ans_rel_list=[],
                ans_event_list=[],
                text=neg_text,
                bg=""
            )
            neg_total_data.append(a_piece_of_neg_data)
    # 按比例划分 neg_total_data
    neg_total_data_len = len(neg_total_data)
    split_ratios = {
        'train': 0.8,
        'test': 0.1,
        'dev': 0.1
    }
    for i in range(neg_total_data_len):
        if i < neg_total_data_len * split_ratios['train']:
            train_dataset.append(neg_total_data[i])
        elif i < neg_total_data_len * (split_ratios['train'] + split_ratios['test']):
            test_dataset.append(neg_total_data[i])
        else:
            dev_dataset.append(neg_total_data[i])
    return train_dataset, test_dataset, dev_dataset

# 打乱数据集
def shuffle_dataset(datasets: list):
    for dataset in datasets:
        random.shuffle(dataset)

# 保存数据集到 jsonl 文件
def save_to_jsonl(data_list, filename):
    def default(obj):
        return obj.__dict__
    try:
        with open(os.path.join(cur_dir, filename), 'w', encoding='utf-8') as f:
            for data in data_list:
                # 将字典转换为JSON格式的字符串，并写入文件，每个对象占一行
                json_line = json.dumps(data.assemble_data(), ensure_ascii=False, default=default) + '\n'
                f.write(json_line)
    except Exception as e:
        print(f"Error occurred while saving data to {filename}: {e}")
# endregion 预处理数据集的一系列方法

# region 初始化
story_en = get_story_from_txt(cur_dir, "../story_en.txt")
story_span_list = get_story_span_list(story_en)
id = ""
ent_list = get_ent_list_from_file(cur_dir, "../ent_list.json")
rel_list = []
event_dict = get_event_dict_from_file(cur_dir, "../event_dict.json")
ans_ent_list = []
ans_rel_list = []
ans_event_list = []
text = ""
bg = ""
# endregion 初始化

# 生成数据集
total_data, train_dataset, test_dataset, dev_dataset = gen_ans_ent_rel_event_text_from_json_file(cur_dir, '../../finetune_data_1.json', story_en, story_span_list)

# 验证数据集
if verify_dataset([train_dataset, test_dataset, dev_dataset]):
    print("数据集验证通过")
else:
    print("数据集验证失败")

# 生成负向数据
train_dataset, test_dataset, dev_dataset = gen_neg_data(total_data, train_dataset, test_dataset, dev_dataset)

# 打乱数据集
shuffle_dataset([train_dataset, test_dataset, dev_dataset])

# 生成数据 id
for i in range(len(train_dataset)):
    train_dataset[i].id = gen_id('train', i+1)
for i in range(len(test_dataset)):
    test_dataset[i].id = gen_id('test', i+1)
for i in range(len(dev_dataset)):
    dev_dataset[i].id = gen_id('dev', i+1)

# 保存数据集
save_to_jsonl(train_dataset, "train.jsonl")
save_to_jsonl(test_dataset, "test.jsonl")
save_to_jsonl(dev_dataset, "dev.jsonl")

# 打印数据集长度
print(f"已保存 {len(train_dataset)} 条训练数据，{len(test_dataset)} 条测试数据，{len(dev_dataset)} 条验证数据")
